{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch version: 2.2.2\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "print(\"torch version:\", version(\"torch\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89], # Your     (x^1)\n",
    "   [0.55, 0.87, 0.66], # journey  (x^2)\n",
    "   [0.57, 0.85, 0.64], # starts   (x^3)\n",
    "   [0.22, 0.58, 0.33], # with     (x^4)\n",
    "   [0.77, 0.25, 0.10], # one      (x^5)\n",
    "   [0.05, 0.80, 0.55]] # step     (x^6)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])\n"
     ]
    }
   ],
   "source": [
    "query = inputs[1]\n",
    "\n",
    "attn_scores_2 = torch.empty(inputs.shape[0])\n",
    "for i,x_i in enumerate(inputs):\n",
    "    attn_scores_2[i] = torch.dot(x_i, query)\n",
    "\n",
    "print(attn_scores_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "atten weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])\n",
      "sum: tensor(1.0000)\n"
     ]
    }
   ],
   "source": [
    "attn_weights_2_tmp = attn_scores_2/attn_scores_2.sum()\n",
    "\n",
    "print(\"atten weights:\", attn_weights_2_tmp)\n",
    "print(\"sum:\", attn_weights_2_tmp.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attention ws: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])\n",
      "sum: tensor(1.)\n"
     ]
    }
   ],
   "source": [
    "def softmax_naive(x):\n",
    "    return torch.exp(x) /torch.exp(x).sum(dim=0)\n",
    "\n",
    "attn_weights_2_naive = softmax_naive(attn_scores_2)\n",
    "\n",
    "print(\"attention ws:\", attn_weights_2_naive)\n",
    "print(\"sum:\", attn_weights_2_naive.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attention ws: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])\n",
      "sum： tensor(1.)\n"
     ]
    }
   ],
   "source": [
    "attn_weights_2 = torch.softmax(attn_scores_2, dim=0)\n",
    "\n",
    "print(\"attention ws:\", attn_weights_2)\n",
    "print(\"sum：\", attn_weights_2.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4419, 0.6515, 0.5683])\n"
     ]
    }
   ],
   "source": [
    "query = inputs[1]\n",
    "\n",
    "context_vec_2 = torch.zeros(query.shape)\n",
    "\n",
    "for i, x_i in enumerate(inputs):\n",
    "    context_vec_2 += attn_weights_2[i]*x_i\n",
    "\n",
    "print(context_vec_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n"
     ]
    }
   ],
   "source": [
    "attn_scores = torch.zeros(6,6)\n",
    "\n",
    "for i,x_i in enumerate(inputs):\n",
    "    for j, x_j in enumerate(inputs):\n",
    "        attn_scores[i,j] = torch.dot(x_i,x_j)\n",
    "print(attn_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n"
     ]
    }
   ],
   "source": [
    "attn_scores = inputs @ inputs.T\n",
    "\n",
    "print(attn_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],\n",
      "        [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],\n",
      "        [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],\n",
      "        [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],\n",
      "        [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],\n",
      "        [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])\n"
     ]
    }
   ],
   "source": [
    "attn_weights = torch.softmax(attn_scores, dim=1)\n",
    "print(attn_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "all row sum: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])\n"
     ]
    }
   ],
   "source": [
    "print(\"all row sum:\", attn_weights.sum(dim=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_context_vecs = attn_weights @ inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.4421, 0.5931, 0.5790],\n",
      "        [0.4419, 0.6515, 0.5683],\n",
      "        [0.4431, 0.6496, 0.5671],\n",
      "        [0.4304, 0.6298, 0.5510],\n",
      "        [0.4671, 0.5910, 0.5266],\n",
      "        [0.4177, 0.6503, 0.5645]])\n"
     ]
    }
   ],
   "source": [
    "print(all_context_vecs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.5500, 0.8700, 0.6600]) 3 2\n"
     ]
    }
   ],
   "source": [
    "x_2 = inputs[1]\n",
    "d_in = inputs.shape[1]\n",
    "d_out = 2\n",
    "\n",
    "print(x_2, d_in, d_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "W_q = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
    "W_k = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
    "W_v = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4306, 1.4551])\n",
      "tensor([0.4433, 1.1419])\n",
      "tensor([0.3951, 1.0037])\n"
     ]
    }
   ],
   "source": [
    "q_2 = x_2@W_q\n",
    "k_2 = x_2@W_k\n",
    "v_2 = x_2@W_v\n",
    "\n",
    "print(q_2)\n",
    "print(k_2)\n",
    "print(v_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.4300, 0.1500, 0.8900],\n",
      "        [0.5500, 0.8700, 0.6600],\n",
      "        [0.5700, 0.8500, 0.6400],\n",
      "        [0.2200, 0.5800, 0.3300],\n",
      "        [0.7700, 0.2500, 0.1000],\n",
      "        [0.0500, 0.8000, 0.5500]])\n",
      "tensor([[0.3669, 0.7646],\n",
      "        [0.4433, 1.1419],\n",
      "        [0.4361, 1.1156],\n",
      "        [0.2408, 0.6706],\n",
      "        [0.1827, 0.3292],\n",
      "        [0.3275, 0.9642]])\n",
      "tensor([[0.1855, 0.8812],\n",
      "        [0.3951, 1.0037],\n",
      "        [0.3879, 0.9831],\n",
      "        [0.2393, 0.5493],\n",
      "        [0.1492, 0.3346],\n",
      "        [0.3221, 0.7863]])\n"
     ]
    }
   ],
   "source": [
    "keys = inputs@W_k\n",
    "values = inputs@W_v\n",
    "\n",
    "print(inputs)\n",
    "print(keys)\n",
    "print(values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4433, 1.1419])\n",
      "tensor([0.4306, 1.4551])\n",
      "tensor(1.8524)\n"
     ]
    }
   ],
   "source": [
    "keys_2 = keys[1]\n",
    "attn_score_22 = q_2.dot(keys_2)\n",
    "print(keys_2)\n",
    "print(q_2)\n",
    "print(attn_score_22)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])\n"
     ]
    }
   ],
   "source": [
    "attn_scores_2 = q_2@keys.T\n",
    "print(attn_scores_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])\n"
     ]
    }
   ],
   "source": [
    "d_k = keys.shape[1]\n",
    "attn_weights_2 = torch.softmax(attn_scores_2/d_k**0.5, dim=-1)\n",
    "print(attn_weights_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.3061, 0.8210])\n"
     ]
    }
   ],
   "source": [
    "context_vec_2 = attn_weights_2@values\n",
    "print(context_vec_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2996, 0.8053],\n",
      "        [0.3061, 0.8210],\n",
      "        [0.3058, 0.8203],\n",
      "        [0.2948, 0.7939],\n",
      "        [0.2927, 0.7891],\n",
      "        [0.2990, 0.8040]], grad_fn=<MmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class SelfAttention_v1(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out):\n",
    "        super().__init__()\n",
    "        self.W_q = nn.Parameter(torch.rand(d_in, d_out))\n",
    "        self.W_k = nn.Parameter(torch.rand(d_in, d_out))\n",
    "        self.W_v = nn.Parameter(torch.rand(d_in, d_out))\n",
    "\n",
    "    def forward(self, x):\n",
    "        keys = x @ self.W_k\n",
    "        queries = x @ self.W_q\n",
    "        values = x @ self.W_v\n",
    "\n",
    "        attn_scores = queries @ keys.T\n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores /keys.shape[-1]**0.5, dim=-1\n",
    "        )\n",
    "\n",
    "        context_vec = attn_weights @ values\n",
    "        return context_vec\n",
    "\n",
    "torch.manual_seed(123)\n",
    "sa_v1 = SelfAttention_v1(d_in, d_out)\n",
    "print(sa_v1(inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.5337, -0.1051],\n",
      "        [-0.5323, -0.1080],\n",
      "        [-0.5323, -0.1079],\n",
      "        [-0.5297, -0.1076],\n",
      "        [-0.5311, -0.1066],\n",
      "        [-0.5299, -0.1081]], grad_fn=<MmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "class SelfAttention_v2(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.W_q = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_k = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_v = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "\n",
    "    def forward(self, x):\n",
    "        keys = self.W_k(x)\n",
    "        queries = self.W_q(x)\n",
    "        values = self.W_v(x)\n",
    "\n",
    "        attn_scores = queries @ keys.T\n",
    "        attn_weights = torch.softmax(attn_scores/keys.shape[-1]**0.5, dim=-1)\n",
    "\n",
    "        context_vec = attn_weights @ values\n",
    "\n",
    "        return context_vec\n",
    "\n",
    "torch.manual_seed(123)\n",
    "sa_v2 = SelfAttention_v2(d_in, d_out)\n",
    "print(sa_v2(inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1717, 0.1762, 0.1761, 0.1555, 0.1627, 0.1579],\n",
      "        [0.1636, 0.1749, 0.1746, 0.1612, 0.1605, 0.1652],\n",
      "        [0.1637, 0.1749, 0.1746, 0.1611, 0.1606, 0.1651],\n",
      "        [0.1636, 0.1704, 0.1702, 0.1652, 0.1632, 0.1674],\n",
      "        [0.1667, 0.1722, 0.1721, 0.1618, 0.1633, 0.1639],\n",
      "        [0.1624, 0.1709, 0.1706, 0.1654, 0.1625, 0.1682]],\n",
      "       grad_fn=<SoftmaxBackward0>)\n"
     ]
    }
   ],
   "source": [
    "# Reuse the query and key weight matrices of th sa_v2 object from the previous section\n",
    "queries = sa_v2.W_q(inputs)\n",
    "keys = sa_v2.W_k(inputs)\n",
    "attn_scores = queries@keys.T\n",
    "\n",
    "attn_weights = torch.softmax(attn_scores/keys.shape[-1]**0.5, dim=-1)\n",
    "print(attn_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 0., 0., 0., 0., 0.],\n",
      "        [1., 1., 0., 0., 0., 0.],\n",
      "        [1., 1., 1., 0., 0., 0.],\n",
      "        [1., 1., 1., 1., 0., 0.],\n",
      "        [1., 1., 1., 1., 1., 0.],\n",
      "        [1., 1., 1., 1., 1., 1.]])\n"
     ]
    }
   ],
   "source": [
    "context_length = attn_scores.shape[0]\n",
    "mask_simple = torch.tril(torch.ones(context_length, context_length))\n",
    "print(mask_simple)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1717, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.1636, 0.1749, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.1637, 0.1749, 0.1746, 0.0000, 0.0000, 0.0000],\n",
      "        [0.1636, 0.1704, 0.1702, 0.1652, 0.0000, 0.0000],\n",
      "        [0.1667, 0.1722, 0.1721, 0.1618, 0.1633, 0.0000],\n",
      "        [0.1624, 0.1709, 0.1706, 0.1654, 0.1625, 0.1682]],\n",
      "       grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "masked_simple = attn_weights*mask_simple\n",
    "print(masked_simple)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.4833, 0.5167, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.3190, 0.3408, 0.3402, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2445, 0.2545, 0.2542, 0.2468, 0.0000, 0.0000],\n",
      "        [0.1994, 0.2060, 0.2058, 0.1935, 0.1953, 0.0000],\n",
      "        [0.1624, 0.1709, 0.1706, 0.1654, 0.1625, 0.1682]],\n",
      "       grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "row_sums = masked_simple.sum(dim=-1, keepdim=True)\n",
    "masked_simple_norm = masked_simple/row_sums\n",
    "print(masked_simple_norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.3111,   -inf,   -inf,   -inf,   -inf,   -inf],\n",
      "        [0.1655, 0.2602,   -inf,   -inf,   -inf,   -inf],\n",
      "        [0.1667, 0.2602, 0.2577,   -inf,   -inf,   -inf],\n",
      "        [0.0510, 0.1080, 0.1064, 0.0643,   -inf,   -inf],\n",
      "        [0.1415, 0.1875, 0.1863, 0.0987, 0.1121,   -inf],\n",
      "        [0.0476, 0.1192, 0.1171, 0.0731, 0.0477, 0.0966]],\n",
      "       grad_fn=<MaskedFillBackward0>)\n"
     ]
    }
   ],
   "source": [
    "mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)\n",
    "masked = attn_scores.masked_fill(mask.bool(), -torch.inf)\n",
    "print(masked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.4833, 0.5167, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.3190, 0.3408, 0.3402, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2445, 0.2545, 0.2542, 0.2468, 0.0000, 0.0000],\n",
      "        [0.1994, 0.2060, 0.2058, 0.1935, 0.1953, 0.0000],\n",
      "        [0.1624, 0.1709, 0.1706, 0.1654, 0.1625, 0.1682]],\n",
      "       grad_fn=<SoftmaxBackward0>)\n"
     ]
    }
   ],
   "source": [
    "attn_weights = torch.softmax(masked/keys.shape[-1]**0.5, dim=-1)\n",
    "print(attn_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[2., 2., 2., 2., 2., 2.],\n",
      "        [0., 2., 0., 0., 0., 0.],\n",
      "        [0., 0., 2., 0., 2., 0.],\n",
      "        [2., 2., 0., 0., 0., 2.],\n",
      "        [2., 0., 0., 0., 0., 2.],\n",
      "        [0., 2., 0., 0., 0., 0.]])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "dropout = torch.nn.Dropout(0.5)\n",
    "example = torch.ones(6,6)\n",
    "\n",
    "print(dropout(example))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 1.0335, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 0.6804, 0.0000, 0.0000, 0.0000],\n",
      "        [0.4889, 0.5090, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.3988, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.3418, 0.0000, 0.0000, 0.0000, 0.0000]],\n",
      "       grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "print(dropout(attn_weights))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 6, 3])\n"
     ]
    }
   ],
   "source": [
    "batch = torch.stack((inputs,inputs), dim=0)\n",
    "print(batch.shape) # 2 inputs with 6 tokens each, and each token has embedding dimention 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.4519,  0.2216],\n",
      "         [-0.5874,  0.0058],\n",
      "         [-0.6300, -0.0632],\n",
      "         [-0.5675, -0.0843],\n",
      "         [-0.5526, -0.0981],\n",
      "         [-0.5299, -0.1081]],\n",
      "\n",
      "        [[-0.4519,  0.2216],\n",
      "         [-0.5874,  0.0058],\n",
      "         [-0.6300, -0.0632],\n",
      "         [-0.5675, -0.0843],\n",
      "         [-0.5526, -0.0981],\n",
      "         [-0.5299, -0.1081]]], grad_fn=<UnsafeViewBackward0>)\n",
      "context_vecs.shape: torch.Size([2, 6, 2])\n"
     ]
    }
   ],
   "source": [
    "class CausalAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.d_out = d_out\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_values = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.dropout = nn.Dropout(dropout) \n",
    "        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length),diagonal=1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, num_tokens, d_in = x.shape\n",
    "        keys = self.W_key(x)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_values(x)\n",
    "\n",
    "        attn_scores = queries @ keys.transpose(1,2)\n",
    "        attn_scores.masked_fill_(\n",
    "            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf\n",
    "        )\n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores/keys.shape[-1]**0.5, dim=-1\n",
    "        )\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        context_vec = attn_weights @ values\n",
    "        return context_vec\n",
    "    \n",
    "torch.manual_seed(123)\n",
    "\n",
    "context_length = batch.shape[1]\n",
    "\n",
    "ca = CausalAttention(d_in, d_out, context_length, 0.0)\n",
    "\n",
    "context_vecs = ca(batch)\n",
    "\n",
    "print(context_vecs)\n",
    "print(\"context_vecs.shape:\", context_vecs.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.4519,  0.2216,  0.4772,  0.1063],\n",
      "         [-0.5874,  0.0058,  0.5891,  0.3257],\n",
      "         [-0.6300, -0.0632,  0.6202,  0.3860],\n",
      "         [-0.5675, -0.0843,  0.5478,  0.3589],\n",
      "         [-0.5526, -0.0981,  0.5321,  0.3428],\n",
      "         [-0.5299, -0.1081,  0.5077,  0.3493]],\n",
      "\n",
      "        [[-0.4519,  0.2216,  0.4772,  0.1063],\n",
      "         [-0.5874,  0.0058,  0.5891,  0.3257],\n",
      "         [-0.6300, -0.0632,  0.6202,  0.3860],\n",
      "         [-0.5675, -0.0843,  0.5478,  0.3589],\n",
      "         [-0.5526, -0.0981,  0.5321,  0.3428],\n",
      "         [-0.5299, -0.1081,  0.5077,  0.3493]]], grad_fn=<CatBackward0>)\n",
      "context_vecs.shape: torch.Size([2, 6, 4])\n"
     ]
    }
   ],
   "source": [
    "class MultiHeadAttentionWrapper(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList(\n",
    "            [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias)\n",
    "             for _ in range(num_heads)]\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.cat([head(x) for head in self.heads], dim=-1)\n",
    "    \n",
    "torch.manual_seed(123)\n",
    "\n",
    "context_length = batch.shape[1] #the number of the tokens\n",
    "\n",
    "d_in, d_out = 3, 2\n",
    "mha = MultiHeadAttentionWrapper(\n",
    "    d_in, d_out, context_length, 0.0, num_heads=2\n",
    ")\n",
    "\n",
    "context_vecs = mha(batch)\n",
    "\n",
    "print(context_vecs)\n",
    "print(\"context_vecs.shape:\", context_vecs.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.2595, 0.4014],\n",
      "         [0.2583, 0.4014],\n",
      "         [0.2583, 0.4014],\n",
      "         [0.2575, 0.4031],\n",
      "         [0.2582, 0.4026],\n",
      "         [0.2575, 0.4028]],\n",
      "\n",
      "        [[0.2595, 0.4014],\n",
      "         [0.2583, 0.4014],\n",
      "         [0.2583, 0.4014],\n",
      "         [0.2575, 0.4031],\n",
      "         [0.2582, 0.4026],\n",
      "         [0.2575, 0.4028]]], grad_fn=<ViewBackward0>)\n",
      "context_vecs.shape: torch.Size([2, 6, 2])\n"
     ]
    }
   ],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        assert (d_out % num_heads ==0), \"d_out must be divisible by num_heads\"\n",
    "        self.d_out = d_out\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = d_out//num_heads  #Reduce the projection dim to match desired output\n",
    "\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.out_proj = nn.Linear(d_out, d_out)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.register_buffer(\n",
    "            \"mask\",\n",
    "            torch.triu(torch.ones(context_length, context_length), diagonal=1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, num_tokens, d_in = x.shape\n",
    "        \n",
    "        keys = self.W_key(x) #shape:(b, num_tokens, d_out)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "\n",
    "        # we implicitly split the matrix by adding a num_heads dimention\n",
    "        # unroll last dim:(b,num_tokens, d_out)->(b,num_tokens,num_heads,head_dim)\n",
    "        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        values = values.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        \n",
    "\n",
    "        # Transpose:(b,num_tokens,num_heads,head_dim)->(b,num_heads,num_tokens,head_dim)\n",
    "        keys = keys.transpose(1,2)\n",
    "        queries = queries.transpose(1,2)\n",
    "        values = values.transpose(1,2)\n",
    "\n",
    "        #compute scaled dot-product attention(aka self-attention) with a causal mask\n",
    "        attn_scores = queries @ keys.transpose(2,3) #dot product for each head\n",
    "\n",
    "        # original mask trancated to the num of tokens and converted to boolean\n",
    "        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]\n",
    "\n",
    "        #use the mask to fill attention scores\n",
    "        attn_scores.masked_fill(mask_bool, -torch.inf)\n",
    "\n",
    "        attn_weights = torch.softmax(attn_scores/keys.shape[-1]**0.5, dim=-1)\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        #shape:(b,num_tokens,num_heads,head_dim)\n",
    "        context_vec = (attn_weights @ values).transpose(1,2)\n",
    "\n",
    "        #combine heads, where self.d_out=self.num_heads*self.head_dim\n",
    "        context_vec = context_vec.contiguous().view(b,num_tokens,self.d_out)\n",
    "        context_vec = self.out_proj(context_vec) #optional projection\n",
    "\n",
    "        return context_vec\n",
    "    \n",
    "torch.manual_seed(123)\n",
    "\n",
    "batch_size, context_length, d_in = batch.shape\n",
    "d_out = 2\n",
    "mha = MultiHeadAttention(d_in,d_out,context_length,0.0,num_heads=2)\n",
    "\n",
    "context_vecs = mha(batch)\n",
    "\n",
    "print(context_vecs)\n",
    "print(\"context_vecs.shape:\", context_vecs.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 2, 3, 4])\n",
      "tensor([[[[0.2745, 0.8993, 0.7179],\n",
      "          [0.6584, 0.0390, 0.7058],\n",
      "          [0.2775, 0.9268, 0.9156],\n",
      "          [0.8573, 0.7388, 0.4340]],\n",
      "\n",
      "         [[0.0772, 0.4066, 0.4606],\n",
      "          [0.3565, 0.2318, 0.5159],\n",
      "          [0.1479, 0.4545, 0.4220],\n",
      "          [0.5331, 0.9737, 0.5786]]]])\n",
      "tensor([[[[1.3208, 1.1631, 1.2879],\n",
      "          [1.1631, 2.2150, 1.8424],\n",
      "          [1.2879, 1.8424, 2.0402]],\n",
      "\n",
      "         [[0.4391, 0.7003, 0.5903],\n",
      "          [0.7003, 1.3737, 1.0620],\n",
      "          [0.5903, 1.0620, 0.9912]]]])\n"
     ]
    }
   ],
   "source": [
    "# (b,num_heads,num_tokens,head_dim) = (1,2,3,4)\n",
    "a = torch.tensor([[[[0.2745, 0.6584, 0.2775, 0.8573],\n",
    "                    [0.8993, 0.0390, 0.9268, 0.7388],\n",
    "                    [0.7179, 0.7058, 0.9156, 0.4340]],\n",
    "\n",
    "                   [[0.0772, 0.3565, 0.1479, 0.5331],\n",
    "                    [0.4066, 0.2318, 0.4545, 0.9737],\n",
    "                    [0.4606, 0.5159, 0.4220, 0.5786]]]])\n",
    "print(a.shape)\n",
    "print(a.transpose(2,3))\n",
    "print(a@a.transpose(2,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2745, 0.6584, 0.2775, 0.8573],\n",
      "        [0.8993, 0.0390, 0.9268, 0.7388],\n",
      "        [0.7179, 0.7058, 0.9156, 0.4340]])\n",
      "-----\n",
      "tensor([[0.2745, 0.8993, 0.7179],\n",
      "        [0.6584, 0.0390, 0.7058],\n",
      "        [0.2775, 0.9268, 0.9156],\n",
      "        [0.8573, 0.7388, 0.4340]])\n",
      "-------\n",
      "tensor([[1.3208, 1.1631, 1.2879],\n",
      "        [1.1631, 2.2150, 1.8424],\n",
      "        [1.2879, 1.8424, 2.0402]])\n",
      "--------\n",
      "tensor([[0.4391, 0.7003, 0.5903],\n",
      "        [0.7003, 1.3737, 1.0620],\n",
      "        [0.5903, 1.0620, 0.9912]])\n"
     ]
    }
   ],
   "source": [
    "first_head = a[0,0,:,:]\n",
    "print(first_head)\n",
    "print(\"-----\")\n",
    "print(first_head.T)\n",
    "first_res = first_head @ first_head.T\n",
    "print(\"-------\")\n",
    "print(first_res)\n",
    "print(\"--------\")\n",
    "second_head = a[0,1,:,:]\n",
    "second_res = second_head@second_head.T\n",
    "print(second_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
